import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append('/root/ptv3/PointTransformerV3/Pointcept/pointcept/utils')
from visualization import *
import open3d as o3d
import numpy as np

S3DIS_COLOR_MAP_20 = {
    0: (233.0, 228.0, 107.0),
    1: (113.0, 186.0, 235.0),
    2: (214.0, 138.0, 96.0),
    3: (240.0, 148.0, 136.0),
    4: (255.0, 61.0, 61.0),
    5: (77.0, 173.0, 83.0),
    6: (108.0, 134.0, 74.0),
    7: (48.0, 58.0, 121.0),
    8: (107.0, 56.0, 113.0),
    9: (96.0, 196.0, 117.0),
    10: (95.0, 95.0, 91.0),
    11: (239.0, 239.0, 239.0),
    12: (96.0, 131.0, 136.0),
}

def base_test():
    coord = np.load("/root/ptv3/PointTransformerV3/Pointcept/data/s3dis/Area_5/conferenceRoom_1/coord.npy")
    # color = np.load("/root/ptv3/PointTransformerV3/Pointcept/data/s3dis/Area_5/conferenceRoom_1/color.npy")
    # color = color.astype(np.float64)/255  # 确保颜色是 float64 类型
    segment = np.load("/root/ptv3/PointTransformerV3/Pointcept/data/s3dis/Area_5/conferenceRoom_1/segment.npy")
    pred = np.load("/root/ptv3/PointTransformerV3/Pointcept/exp/s3dis/semseg-pt-v3m1-0-base/result/Area_5-conferenceRoom_1_pred.npy")
    # instance = np.load("/root/ptv3/PointTransformerV3/Pointcept/data/s3dis/Area_5/conferenceRoom_1/instance.npy")
    color_seg = np.array([S3DIS_COLOR_MAP_20[i.item()] for i in segment], dtype=np.float64) / 255.0  # 确保颜色是 float64 类型
    color_pred = np.array([S3DIS_COLOR_MAP_20[i.item()] for i in pred], dtype=np.float64) / 255.0  # 确保颜色是 float64 类型
    # print(color.shape, color.dtype)  # 应输出类似 (N,3) float64
    # print(set(segment.flatten()))  # 打印前10个点的颜色

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(coord)
    pcd.colors = o3d.utility.Vector3dVector(color_pred)
    o3d.io.write_point_cloud("pred.ply", pcd)

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(coord)
    pcd.colors = o3d.utility.Vector3dVector(color_seg)
    o3d.io.write_point_cloud("seg.ply", pcd)


if __name__ == "__main__":
    # save_point_cloud(coord, color, "test.ply")
    base_test()
    print("Point cloud saved successfully.")